use socket2::Socket;
use std::error::Error;
use tokio::{
    io::{split, AsyncReadExt, AsyncWriteExt},
    net::TcpListener,
    net::TcpStream,
    time::{sleep, timeout, Duration},
};
use tokio_socks::tcp::Socks5Stream;
use url;

// 上游代理配置
const PROXY_URL: &str = "socks5://bcfe09c4ae79075a11c1__cr.us:5759f3e22c7b725a@144.76.124.83:824";

// 超时设置
const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
const READ_TIMEOUT: Duration = Duration::from_secs(60);
const WRITE_TIMEOUT: Duration = Duration::from_secs(60);
const RETRY_DELAY: Duration = Duration::from_millis(100);
const MAX_RETRIES: u32 = 5;

type BoxError = Box<dyn Error + Send + Sync>;

// 解析代理URL的函数
fn parse_proxy_url() -> Result<(String, String, String), BoxError> {
    let url = url::Url::parse(PROXY_URL)?;
    Ok((
        format!(
            "{}:{}",
            url.host_str().unwrap_or(""),
            url.port().unwrap_or(1080)
        ),
        url.username().to_string(),
        url.password().unwrap_or("").to_string(),
    ))
}

#[tokio::main]
async fn main() -> Result<(), BoxError> {
    let listener = TcpListener::bind("127.0.0.1:1080").await?;
    println!("SOCKS5h proxy listening on: 127.0.0.1:1080");

    loop {
        let (client, addr) = listener.accept().await?;
        println!("New connection from: {}", addr);

        tokio::spawn(async move {
            if let Err(e) = handle_client_with_timeout(client).await {
                eprintln!("Error handling client {}: {}", addr, e);
            }
        });
    }
}

async fn handle_client_with_timeout(client: TcpStream) -> Result<(), BoxError> {
    let peer_addr = client.peer_addr()?;
    println!("Starting handling client: {}", peer_addr);
    match timeout(CONNECT_TIMEOUT, handle_client(client)).await {
        Ok(result) => {
            if let Err(e) = result {
                println!("Client {} error: {}", peer_addr, e);
                return Err(e);
            }
            println!("Client {} completed successfully", peer_addr);
            Ok(())
        }
        Err(e) => {
            println!("Client {} timeout: {}", peer_addr, e);
            Err(e.into())
        }
    }
}

async fn handle_client(mut client: TcpStream) -> Result<(), BoxError> {
    let peer_addr = client.peer_addr()?;
    client.set_nodelay(true)?;

    // 添加TCP保活机制
    let socket = Socket::from(client.into_std()?);
    socket.set_keepalive(true)?;
    client = TcpStream::from_std(socket.into())?;

    // 1. 处理握手请求
    let mut buf = [0u8; 4096];
    let n = timeout(READ_TIMEOUT, client.read(&mut buf)).await??;
    println!("Client {} handshake received {} bytes", peer_addr, n);
    if n < 2 || buf[0] != 0x05 {
        return Err("Not a SOCKS5 request".into());
    }

    timeout(WRITE_TIMEOUT, client.write_all(&[0x05, 0x00])).await??;
    println!("Client {} handshake completed", peer_addr);

    // 2. 处理连接请求
    let n = timeout(READ_TIMEOUT, client.read(&mut buf)).await??;
    if n < 5 || buf[0] != 0x05 {
        return Err("Invalid SOCKS5 request".into());
    }

    let cmd = buf[1];
    if cmd != 0x01 {
        return Err("Only CONNECT method is supported".into());
    }

    // 解析目标地址和端口
    let addr_type = buf[3];
    let (target_addr, target_port, _) = match addr_type {
        0x01 => {
            // IPv4
            let addr = format!("{}.{}.{}.{}", buf[4], buf[5], buf[6], buf[7]);
            let port = ((buf[8] as u16) << 8) | (buf[9] as u16);
            (addr, port, 10)
        }
        0x03 => {
            // Domain
            let len = buf[4] as usize;
            let domain = String::from_utf8_lossy(&buf[5..5 + len]).to_string();
            let port = ((buf[5 + len] as u16) << 8) | (buf[5 + len + 1] as u16);
            (domain, port, 5 + len + 2)
        }
        _ => return Err("Unsupported address type".into()),
    };

    let target = format!("{}:{}", target_addr, target_port);
    let (proxy_host, proxy_username, proxy_password) = parse_proxy_url()?;

    // 3. 连接上游代理
    let mut upstream = None;
    let mut retries = 0;

    while retries < MAX_RETRIES {
        println!(
            "Attempting to connect to proxy (attempt {}/{})",
            retries + 1,
            MAX_RETRIES
        );

        match Socks5Stream::connect_with_password(
            proxy_host.as_str(),
            target.as_str(),
            &proxy_username,
            &proxy_password,
        )
        .await
        {
            Ok(stream) => {
                println!(
                    "Successfully connected to proxy after {} attempts",
                    retries + 1
                );
                upstream = Some(stream);
                break;
            }
            Err(e) => {
                eprintln!("Connection attempt {} failed: {}", retries + 1, e);
                retries += 1;
                if retries < MAX_RETRIES {
                    sleep(RETRY_DELAY).await;
                }
            }
        }
    }

    let upstream = upstream.ok_or("Failed to connect after all retries")?;

    // 4. 发送成功响应给客户端
    let response = [0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
    timeout(WRITE_TIMEOUT, client.write_all(&response)).await??;

    // 5. 开始数据传输
    println!(
        "Client {} connection established, starting bidirectional transfer",
        peer_addr
    );
    transfer_data_with_timeout(client, upstream.into_inner()).await?;
    println!("Client {} transfer completed", peer_addr);

    Ok(())
}

async fn transfer_data_with_timeout(
    client: TcpStream,
    upstream: TcpStream,
) -> Result<(), BoxError> {
    let peer_addr = client.peer_addr()?;
    println!("Starting data transfer for client: {}", peer_addr);

    // 在分割之前设置 keepalive
    let socket = Socket::from(client.into_std()?);
    socket.set_keepalive(true)?;
    let client = TcpStream::from_std(socket.into())?;

    let socket = Socket::from(upstream.into_std()?);
    socket.set_keepalive(true)?;
    let upstream = TcpStream::from_std(socket.into())?;

    let (client_read, client_write) = split(client);
    let (upstream_read, upstream_write) = split(upstream);

    let client_to_upstream = tokio::spawn(async move {
        let mut buf = [0u8; 64 * 1024]; // 增加缓冲区大小到64KB
        let mut read_half = client_read;
        let mut write_half = upstream_write;
        let mut total_bytes = 0u64;

        loop {
            match timeout(READ_TIMEOUT, read_half.read(&mut buf)).await {
                Ok(Ok(0)) => {
                    println!(
                        "Client closed connection, total bytes sent: {}",
                        total_bytes
                    );
                    break;
                }
                Ok(Ok(n)) => {
                    if let Err(e) = timeout(WRITE_TIMEOUT, write_half.write_all(&buf[..n])).await {
                        eprintln!("Failed to write to upstream: {}", e);
                        break;
                    }
                    total_bytes += n as u64;
                }
                Err(e) => {
                    eprintln!("Timeout reading from client: {}", e);
                    break;
                }
                Ok(Err(e)) => {
                    eprintln!("Failed to read from client: {}", e);
                    break;
                }
            }
        }
    });

    let upstream_to_client = tokio::spawn(async move {
        let mut buf = [0u8; 64 * 1024]; // 增加缓冲区大小到64KB
        let mut read_half = upstream_read;
        let mut write_half = client_write;
        let mut total_bytes = 0u64;

        loop {
            match timeout(READ_TIMEOUT, read_half.read(&mut buf)).await {
                Ok(Ok(0)) => {
                    println!(
                        "Upstream closed connection, total bytes received: {}",
                        total_bytes
                    );
                    break;
                }
                Ok(Ok(n)) => {
                    if let Err(e) = timeout(WRITE_TIMEOUT, write_half.write_all(&buf[..n])).await {
                        eprintln!("Failed to write to client: {}", e);
                        break;
                    }
                    total_bytes += n as u64;
                }
                Err(e) => {
                    eprintln!("Timeout reading from upstream: {}", e);
                    break;
                }
                Ok(Err(e)) => {
                    eprintln!("Failed to read from upstream: {}", e);
                    break;
                }
            }
        }
    });

    tokio::select! {
        _ = client_to_upstream => println!("Client {} to upstream transfer completed", peer_addr),
        _ = upstream_to_client => println!("Upstream to client {} transfer completed", peer_addr),
    }

    Ok(())
}
